import java.io.*;
import java.net.*;
import java.nio.ByteBuffer;
import java.util.*;

public class DNS {

    private static final int DNS_PORT = 53;
    private static final int UDP_MAX_SIZE = 512;
    private static final String ROOT_SERVER = "192.168.99.81"; // 根域名服务器

    public static void main(String[] args) {
        try {
            // 启动 UDP 服务器
            new Thread(() -> startUdpServer()).start();
            // 启动 TCP 服务器
            new Thread(() -> startTcpServer()).start();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 启动 UDP 服务器
    private static void startUdpServer() {
        try (DatagramSocket socket = new DatagramSocket(DNS_PORT)) {
            System.out.println("UDP DNS Server started on port " + DNS_PORT);
            while (true) {
                byte[] buffer = new byte[UDP_MAX_SIZE];
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
                socket.receive(packet);
                // 处理 DNS 请求
                System.out.println("UDP===");
                byte[] response = resolve(packet.getData());
                if (response != null) {
                    DatagramPacket responsePacket = new DatagramPacket(response, response.length, packet.getAddress(), packet.getPort());
                    socket.send(responsePacket);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 启动 TCP 服务器
    private static void startTcpServer() {
        try (ServerSocket serverSocket = new ServerSocket(DNS_PORT)) {
            System.out.println("TCP DNS Server started on port " + DNS_PORT);
            while (true) {
                try (Socket clientSocket = serverSocket.accept();
                    InputStream in = clientSocket.getInputStream();
                     OutputStream out = clientSocket.getOutputStream()) {
                    // 读取请求长度
                    byte[] lengthBytes = new byte[2];
                    in.read(lengthBytes);
                    int length = ByteBuffer.wrap(lengthBytes).getShort() & 0xFFFF;

                    // 读取 DNS 请求
                    byte[] request = new byte[length];
                    in.read(request);

                    // 处理 DNS 请求
                    System.out.println("TCP===");
                    byte[] response = resolve(request);

                    if (response != null) {
                        // 发送响应长度
                        out.write(ByteBuffer.allocate(2).putShort((short) response.length).array());
                        // 发送响应内容
                        out.write(response);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    // 递归解析 DNS 查询
    private static byte[] resolve(byte[] query) {
        try {
            // 从根域名服务器开始查询
            String currentServer = ROOT_SERVER;
            byte[] response = null;
            while (true) {
                // 发送查询请求
                try (DatagramSocket socket = new DatagramSocket()) {
                    socket.setSoTimeout(200); // 设置超时时间
                    // 发送 UDP 查询
                    DatagramPacket packet = new DatagramPacket(query, query.length, InetAddress.getByName(currentServer), DNS_PORT);
                    socket.send(packet);
                    // 接收响应
                    byte[] buffer = new byte[UDP_MAX_SIZE];
                    DatagramPacket responsePacket = new DatagramPacket(buffer, buffer.length);
                    socket.receive(responsePacket);
                    // 解析响应
                    response =  responsePacket.getData();
                } catch (Exception e) {
                    e.printStackTrace();
                }
                // 检查响应
                if (response != null && response.length > 0) {
                    // 如果响应中包含答案，返回结果
                    return response;
                } else {
                    // 没有找到结果
                    throw new RuntimeException("DNS resolution failed");
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}

